'''
Target: Python AndroidToolBox Fastboot Flash Module
Author: LZR@BUAA
Date:   08/26/2020
'''
import os
import re
import sys
import time

from ADBMethods import *
from SysMethods import *

TAR = []
path = os.getcwd()
list_dir = os.listdir(path+"\\images\\")
num_of_TARs = 0


def findTarPacks():
    global num_of_TARs
    for i in list_dir:
        if i.split(".")[-1] == "tgz" or i.split(".")[-1] == "gz":
            TAR.append(i)
            del i
            num_of_TARs += 1


'''
Reboot
'''


def rebootList():
        op = input("请选择操作：\n[1]重启到系统\n[2]重启到Recovery\n[3]重启到BootLoader\n[0]退出程序")
        if op == '1':
            rebootUI()
        elif op == '2':
            rebootRec()
        elif op == '3':
            rebootBL()
        elif op == '0':
            exitProgram(1)
        exitProgram(0)


'''
Flash
'''


def directFlash(mode, file):
    print(cmd("fastboot flash "+mode+" "+file))


def fastbootWipe():
    cmd("fastboot erase userdata")
    cmd("fastboot erase cache")


def Flashmain():
    os.system("COLOR A")
    print("########################################")
    print("")
    print("            FastBoot烧写工具")
    print("              By:LZR@BUAA")
    print("")
    print("########################################")
    print("欢迎使用!\n")
    print("****************风险提示******************")
    print("不正确操作可能导致设备无法启动甚至损坏，请三思而后行")
    print("[1]卡刷(使用独立Zip/IMG安装)")
    print("[2]线刷(使用官方原版TAR安装)")
    print("[3]重启选项")
    print("[4]退出工具")
    op = input("请选择工作模式 按Enter确认")
    if op == '1':
        print("是否已经刷入第三方Recovery？(Y/N)")
        if InputJudge(2):
            input("准备重启到recovery模式，请在手机端保存工作，按Enter继续……")
            rebootRec()
            while True:
                if cmd("adb devices").find("recovery") != -1:
                    break
                time.sleep(3)
            input("请手动在Recovery高级选项中开启ADB Sideload/ADB 线刷功能 按Enter继续……")
            print("请将"+MobileInfo['brand']+" " +
                  MobileInfo['model'] + " 的zip刷机文件拖至窗口中，并按Enter键……")
            imagefile = input()
            print("这是最后一次提示，确实要将"+imagefile.strip("\\")
                  [-1]+"刷入"+MobileInfo['brand']+" " + MobileInfo['model'] + "吗？(Y/N)")
            if InputJudge(2):
                print(cmd("adb sideload "+imagefile))
        else:
            print("是否要刷入第三方Recovery？(Y/N)")
            if InputJudge(2):
                print("请将该机型的recovery拖入该窗口或输入recovery.TAR的完整路径")
                recdir = input()
                if recdir.find("TAR") == -1:
                    print("该文件不是一个有效的Recovery镜像！")
                rebootBL()
                # 等待重启完毕
                while True:
                    time.sleep(1)
                    if BLConnected():
                        break
                directFlash("recovery", recdir)

            else:

                print("是否使用Fastboot模式进行刷机？(必须已解BL锁)(Y/N)")
                if InputJudge(2):
                    print("#####################高风险功能提醒#####################")
                    print("警告！Fastboot刷机模式风险极高，请自行承担设备损坏等后果")
                    print("除非知道自己在做什么，否则请自行关闭程序！")
                    print("仍然要继续吗？(Y/N)")
                    if InputJudge(2):
                        rebootBL()
                        print("选择模式：\n[Y]批量刷入\n[N]逐个刷入")
                        if InputJudge(2):
                            print(
                                "请将.img或.bin格式刷机包放入打开的目录下")
                            os.startfile(".\\images\\FastbootImages")
                            input("不要修改文件名 按Enter继续")
                            IMG_list = os.listdir(
                                path+"\\images\\FastbootImages\\")
                            print("文件读取中……")
                            for img in IMG_list:
                                if os.path.splitext(img)[1] not in [".img", ".IMG", ".Img", ".bin", ".Bin", ".BIN"]:
                                    IMG_list.remove(img)
                            print("待刷入的文件列表：")
                            # for i in range(len(IMG_list)):
                            # print("["+ i +"]"+IMG_list[i])
                            print(IMG_list)
                            print("请务必保持手机与电脑连接，这是最后一次确认，是否要刷入以上镜像(Y/N)")
                            if InputJudge(2):
                                print("")
                                for img in IMG_list:
                                    print("正在刷入文件："+img)
                                    # print(path+"\\images\\FastbootImages\\"+img)
                                    directFlash(os.path.splitext(img)[
                                                0], path+"\\images\\FastbootImages\\"+img)
                                fastbootWipe()
                            else:
                                exitProgram(1)
                        else:
                            fdir = dragtoWindowGetName()
                            directFlash(os.path.splitext(
                                fdir.split("\\")[-1])[0], fdir)

                exitProgram(1)

    if op == '2':
        print("请将*.tar刷机包放入\images\目录下")
        os.startfile(path + "\\images")
        input("复制完成后请按Enter继续")
        findTarPacks()
        print("正在读取并解压压缩包，约30秒……请稍候")
        for i in range(num_of_TARs):
            unziptar(path + "\\images\\" + TAR[i], path+"\\images\\")
            op = input(
                "请选择工作模式：\n[1]全部擦写并清空数据\n[2]全部擦写但保留内部存储(\sdcard)\n[3]全部擦写并加BoolLoader锁(慎用)")
            if op not in ['1', '2', '3']:
                print("输入有误，程序退出中……")
                exit()
            elif op == '1':
                cmd(path + "\\images\\" + list_dir[i] + "\\"+"flash_all.bat")
            elif op == '2':
                cmd(path + "\\images\\" +
                    list_dir[i] + "\\"+"flash_all_except_storage.bat")
            elif op == '3':
                cmd(path + "\\images\\" +
                    list_dir[i] + "\\"+"flash_all_lock.bat")
        exitProgram(0)

    if op == '3':
        rebootList()


if __name__ == "__main__":
    Flashmain()
